{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import os\n",
    "import re\n",
    "from PIL import Image\n",
    "from dispnetwork import DispNet\n",
    "from tensorflow.keras.preprocessing.image import ImageDataGenerator\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "tf.config.experimental.set_memory_growth(gpus[0], True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 800 images belonging to 1 classes.\n",
      "Found 800 images belonging to 1 classes.\n",
      "Found 800 images belonging to 1 classes.\n",
      "(8, 256, 512, 2) (8, 128, 256)\n"
     ]
    }
   ],
   "source": [
    "###########################读取driving数据\n",
    "path = \"F:/dataset/driving/frames_cleanpass/15mm_focallength/scene_forwards/slow\"\n",
    "pathd = \"F:/dataset/driving/disparity_img/15mm_focallength/scene_forwards/slow\"\n",
    "\n",
    "def trainGenerator(batch_size):\n",
    "\n",
    "    aug_dict = dict(horizontal_flip=True,\n",
    "                        fill_mode='nearest')\n",
    "    \n",
    "    left_datagen = ImageDataGenerator(**aug_dict)\n",
    "    right_datagen = ImageDataGenerator(**aug_dict)\n",
    "    deep_datagen = ImageDataGenerator(**aug_dict)\n",
    "    left_image = left_datagen.flow_from_directory(path,classes=['left'],color_mode = \"grayscale\",\n",
    "        target_size = (256, 512),class_mode = None,batch_size = batch_size, seed=3)\n",
    "    \n",
    "    right_image = right_datagen.flow_from_directory(path,classes=['right'],color_mode = \"grayscale\",\n",
    "        target_size = (256, 512),class_mode = None,batch_size = batch_size, seed=3)\n",
    "    \n",
    "    depth_image = deep_datagen.flow_from_directory(pathd,classes=['left'],color_mode = \"grayscale\",\n",
    "        target_size = (128, 256),class_mode = None,batch_size = batch_size, seed=3)\n",
    "\n",
    "    train_generator = zip(left_image, right_image, depth_image)\n",
    "    for (left,right,deep) in train_generator:\n",
    "        imgl = left/255.\n",
    "        imgr = right/255.  \n",
    "        img = tf.concat([imgl,imgr],axis = 3)\n",
    "        deep = tf.squeeze(deep/255.)\n",
    "        yield (img,deep)\n",
    "        \n",
    "trainset = trainGenerator(batch_size=8)\n",
    "images,deep_images = next(trainset)\n",
    "print(images.shape,deep_images.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def SSIM(x, y):\n",
    "    C1 = 0.01 ** 2\n",
    "    C2 = 0.03 ** 2\n",
    "\n",
    "    mu_x = tf.nn.avg_pool2d(x, 3, 1, 'VALID')\n",
    "    mu_y = tf.nn.avg_pool2d(y, 3, 1, 'VALID')\n",
    "\n",
    "    sigma_x  = tf.nn.avg_pool2d(x ** 2, 3, 1, 'VALID') - mu_x ** 2\n",
    "    sigma_y  = tf.nn.avg_pool2d(y ** 2, 3, 1, 'VALID') - mu_y ** 2\n",
    "    sigma_xy = tf.nn.avg_pool2d(x * y , 3, 1, 'VALID') - mu_x * mu_y\n",
    "\n",
    "    SSIM_n = (2 * mu_x * mu_y + C1) * (2 * sigma_xy + C2)\n",
    "    SSIM_d = (mu_x ** 2 + mu_y ** 2 + C1) * (sigma_x + sigma_y + C2)\n",
    "    SSIM = SSIM_n / SSIM_d\n",
    "    return tf.clip_by_value((1 - SSIM) / 2, 0, 1)\n",
    "\n",
    "def do_primid(inp):\n",
    "    p1 = tf.nn.avg_pool2d(tf.expand_dims(inp,axis = 3),ksize = [1,2,2,1],strides=[1,2,2,1],padding = 'SAME')\n",
    "    p2 = tf.nn.avg_pool2d(tf.expand_dims(inp,axis = 3),ksize = [1,4,4,1],strides=[1,4,4,1],padding = 'SAME')\n",
    "    p3 = tf.nn.avg_pool2d(tf.expand_dims(inp,axis = 3),ksize = [1,8,8,1],strides=[1,8,8,1],padding = 'SAME')\n",
    "    p4 = tf.nn.avg_pool2d(tf.expand_dims(inp,axis = 3),ksize = [1,16,16,1],strides=[1,16,16,1],padding = 'SAME')\n",
    "    return tf.expand_dims(inp,axis=3),p1,p2,p3,p4 \n",
    "\n",
    "def get_loss(pred1,pred2,pred3,pred4,pred5,label):\n",
    "    label1,label2,label3,label4,label5 = do_primid(label)\n",
    "    loss1 = tf.reduce_mean(tf.square(label1 - pred1))\n",
    "    loss2 = tf.reduce_mean(tf.square(label2 - pred2))\n",
    "    loss3 = tf.reduce_mean(tf.square(label3 - pred3))\n",
    "    loss4 = tf.reduce_mean(tf.square(label4 - pred4))\n",
    "    loss5 = tf.reduce_mean(tf.square(label5 - pred5))\n",
    "    return (loss1/2 + loss2/4 + loss3/8 + loss4/16+ loss5/32)\n",
    "\n",
    "def get_ssim_loss(pred1,pred2,pred3,pred4,pred5,label):\n",
    "    label1,label2,label3,label4,label5 = do_primid(label)\n",
    "    loss1 = SSIM(label1,pred1)\n",
    "    loss2 = SSIM(label2,pred2)\n",
    "    loss3 = SSIM(label3,pred3)\n",
    "    loss4 = SSIM(label4,pred4)\n",
    "    loss5 = SSIM(label5,pred5)\n",
    "    loss1 = tf.reduce_mean(loss1)\n",
    "    loss2 = tf.reduce_mean(loss2)\n",
    "    loss3 = tf.reduce_mean(loss3)\n",
    "    loss4 = tf.reduce_mean(loss4)\n",
    "    loss5 = tf.reduce_mean(loss5)\n",
    "#     return loss1\n",
    "    return (loss1/2 + loss2/4 + loss3/8 + loss4/16+ loss5/32)\n",
    "\n",
    "model = DispNet()\n",
    "model.build(input_shape=(None, 512, 1024, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 0 loss: 0.2229122668504715\n",
      "step: 1 loss: 0.22323551774024963\n",
      "step: 2 loss: 0.20008639991283417\n",
      "step: 3 loss: 0.2055431306362152\n",
      "step: 4 loss: 0.17913681268692017\n",
      "step: 5 loss: 0.17968159914016724\n",
      "step: 6 loss: 0.2019653618335724\n",
      "step: 7 loss: 0.1930008977651596\n",
      "step: 8 loss: 0.18112045526504517\n",
      "step: 9 loss: 0.17624476552009583\n",
      "step: 10 loss: 0.17230936884880066\n",
      "step: 11 loss: 0.18626263737678528\n",
      "step: 12 loss: 0.18450568616390228\n",
      "step: 13 loss: 0.17725101113319397\n",
      "step: 14 loss: 0.15937726199626923\n",
      "step: 15 loss: 0.15869247913360596\n",
      "step: 16 loss: 0.1545904576778412\n",
      "step: 17 loss: 0.1556800752878189\n",
      "step: 18 loss: 0.13639438152313232\n",
      "step: 19 loss: 0.1424216628074646\n",
      "step: 20 loss: 0.1465042382478714\n",
      "step: 21 loss: 0.1436818540096283\n",
      "step: 22 loss: 0.1464337408542633\n",
      "step: 23 loss: 0.13816407322883606\n",
      "step: 24 loss: 0.14084424078464508\n",
      "step: 25 loss: 0.1413574516773224\n",
      "step: 26 loss: 0.14470604062080383\n",
      "step: 27 loss: 0.1318208873271942\n",
      "step: 28 loss: 0.1328638643026352\n",
      "step: 29 loss: 0.13123315572738647\n",
      "step: 30 loss: 0.14602473378181458\n",
      "step: 31 loss: 0.11965623497962952\n",
      "step: 32 loss: 0.12718558311462402\n",
      "step: 33 loss: 0.13279442489147186\n",
      "step: 34 loss: 0.11903659254312515\n",
      "step: 35 loss: 0.13233669102191925\n",
      "step: 36 loss: 0.13737139105796814\n",
      "step: 37 loss: 0.1281646490097046\n",
      "step: 38 loss: 0.13288554549217224\n",
      "step: 39 loss: 0.1349676251411438\n",
      "step: 40 loss: 0.12773318588733673\n",
      "step: 41 loss: 0.1171271800994873\n",
      "step: 42 loss: 0.11682331562042236\n",
      "step: 43 loss: 0.1146806851029396\n",
      "step: 44 loss: 0.13026535511016846\n",
      "step: 45 loss: 0.10720469802618027\n",
      "step: 46 loss: 0.13306647539138794\n",
      "step: 47 loss: 0.12137341499328613\n",
      "step: 48 loss: 0.10415231436491013\n",
      "step: 49 loss: 0.11959318071603775\n",
      "step: 50 loss: 0.12038631737232208\n",
      "step: 51 loss: 0.1020176112651825\n",
      "step: 52 loss: 0.10726196318864822\n",
      "step: 53 loss: 0.10114822536706924\n",
      "step: 54 loss: 0.1243632510304451\n",
      "step: 55 loss: 0.10898707062005997\n",
      "step: 56 loss: 0.10878913849592209\n",
      "step: 57 loss: 0.09959004074335098\n",
      "step: 58 loss: 0.10314223915338516\n",
      "step: 59 loss: 0.13047659397125244\n",
      "step: 60 loss: 0.12420233339071274\n",
      "step: 61 loss: 0.11188016831874847\n",
      "step: 62 loss: 0.1032743752002716\n",
      "step: 63 loss: 0.11348416656255722\n",
      "step: 64 loss: 0.12459488958120346\n",
      "step: 65 loss: 0.11276953667402267\n",
      "step: 66 loss: 0.11963380873203278\n"
     ]
    }
   ],
   "source": [
    "\n",
    "model.load_weights(\"./weight\")\n",
    "optimizer = tf.keras.optimizers.Adam(lr = 3e-5)\n",
    "for epoch in range(1):\n",
    "    for step in range(100):\n",
    "        with tf.GradientTape() as tape:\n",
    "            images,deep_images = next(trainset)\n",
    "            pred1,pred2,pred3,pred4,pred5 = model(images)\n",
    "            loss =get_ssim_loss(pred1,pred2,pred3,pred4,pred5,deep_images)\n",
    "            gradients = tape.gradient(loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "            print(\"step:\",step,\"loss:\",float(loss))\n",
    "    model.save_weights(\"./weight\")       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model.load_weights(\"./weight\")\n",
    "images,deep_images = next(trainset)\n",
    "pred1,pred2,pred3,pred4,pred5 = model(images)\n",
    "print(pred1.shape)\n",
    "out = tf.squeeze(pred1)\n",
    "for i in range(8):\n",
    "    fig = plt.figure(figsize = (32,128))\n",
    "    plt.imshow(tf.concat([deep_images[i],out[i]],axis = 1),cmap=\"hot\")\n",
    "    plt.show()\n",
    "    print(\"                         真实↑                                                   预测↑\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
